/*
 * Javlov - a Java toolkit for reinforcement learning with multi-agent support.
 * 
 * Copyright (c) 2009 Matthijs Snel
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov.world.grid;

import java.awt.Shape;
import java.awt.geom.Point2D;
import java.awt.geom.Rectangle2D;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import net.javlov.Action;
import net.javlov.Agent;
import net.javlov.RewardFunction;
import net.javlov.State;
import net.javlov.world.AgentBody;
import net.javlov.world.Body;
import net.javlov.world.CollisionEvent;
import net.javlov.world.CollisionListener;
import net.javlov.world.World;
import net.javlov.world.phys2d.Phys2DBody;
import net.phys2d.math.ROVector2f;

public class GridWorld implements World.Discrete, IGridWorld {

	/**
	 * Maps agents to their bodies. Done this way because an agent is not allowed to have access
	 * to its own environment/body.
	 */
	protected Map<Agent, AgentBody> agentBodyMap;
	
	/**
	 * All the bodies in the world.
	 */
	protected List<Body> bodies;
	
	/**
	 * The grid.
	 */
	protected Grid grid;
	
	private boolean episodeEnd;
	
	private List<CollisionListener> listeners;
	
	protected RewardFunction reward;


	public GridWorld(int width, int height, double cellWidth, double cellHeight) {
		agentBodyMap = new HashMap<Agent, AgentBody>();
		bodies = new ArrayList<Body>();
		grid = new Grid(width, height, cellWidth, cellHeight);
		listeners = new ArrayList<CollisionListener>();
	}
	
	@Override
	public boolean add(Agent a, AgentBody b) {
		if ( addBody(b) ) {
			agentBodyMap.put(a, b);
			return true;
		}
		return false;
	}

	@Override
	public boolean addBody(Body b) {
		if ( bodies.add(b) ) {
			addToCells(b);
			return true;
		}
		return false;
	}

	public void addCollisionListener(CollisionListener listener) {
		listeners.add(listener);
	}
	
	@Override
	public Body getAgentBody(Agent a) {
		return agentBodyMap.get(a);
	}

	@Override
	public double getHeight() {
		return grid.getHeight()*grid.getCellHeight();
	}

	@Override
	public List<Body> getIntersectingObjects(Shape s) {
		// TODO Auto-generated method stub
		return null;
	}

	@Override
	public List<Body> getObjects() {
		return bodies;
	}

	@Override
	public double getTimeStep() {
		// TODO Auto-generated method stub
		return 0;
	}

	@Override
	public double getWidth() {
		return grid.getWidth()*grid.getCellWidth();
	}

	@Override
	public boolean intersectsObject(Shape s) {
		// TODO Auto-generated method stub
		return false;
	}

	@Override
	public boolean remove(Agent a) {
		Body b = agentBodyMap.get(a);
		if ( removeBody(b) )
			return (agentBodyMap.remove(a) == null ? false : true);
		return false;
	}

	@Override
	public boolean removeBody(Body b) {
		if ( bodies.remove(b) ) {
			removeFromCells(b);
			return true;
		}
		return false;
	}

	@Override
	public double executeAction(Action act, Agent a) {
		reward.preAction(act, a);
		act.execute(a);
		return reward.calculateReward();
	}

	/**
	 * Returns the state by calling {@link AgentBody#getState(Agent)} on the agent's body.
	 */
	@Override
	public State getState(Agent a) {
		State s = agentBodyMap.get(a).getState(a);
		s.setTerminal(episodeEnd);
		return s;
	}

	/**
	 * Returns the state dim as indicated by {@link AgentBody#getStateDim()}
	 */
	@Override
	public int getStateDim() {
		Iterator<AgentBody> it = agentBodyMap.values().iterator();
		return it.next().getStateDim();
	}

	@Override
	public void init() {
		episodeEnd = false;	
	}

	@Override
	public void reset() {
		episodeEnd = false;
		//TODO hack for now, just to test for 1 agent
		Body b = agentBodyMap.values().iterator().next();
		GridCell currCell = grid.getCell(b.getX(), b.getY());
		currCell.removeBody(b);
		GridCell startCell = grid.getCell(grid.getCellWidth()/2d, grid.getCellHeight()/2d);
		startCell.addBody(b);
		b.setLocation(startCell.getCenterX(), startCell.getCenterY());	
	}

	@Override
	public boolean rotateBody(Body b, int angle) {
		// TODO Auto-generated method stub
		return true;
	}

	/**
	 * Doesn't do anything. After initialisation of the world, the grid cannot be changed.
	 */
	@Override
	public void setHeight(int height) {}

	/**
	 * Doesn't do anything. After initialisation of the world, the grid cannot be changed.
	 */
	@Override
	public void setWidth(int width) {}

	@Override
	public boolean translateBody(Body b, int dx, int dy) {
		int absx = Math.abs(dx),
			absy = Math.abs(dy);
		if ( absx > 0 && absy > 0 && absx != absy )
			throw new IllegalArgumentException("GridWorld: can only move in straight "
					+ "line, or in diagonal such that dx=dy.");
		
		//now move in unit steps to see if there is anything in the body's path
		int unitdx = (absx == 0 ? 0 : dx/absx),
			unitdy = (absy == 0 ? 0 : dy/absy),
			steps = (absx+absy) / (unitdx+unitdy);
		return translateBody(b, Direction.get(unitdx, unitdy), steps);		
	}
	
	public boolean translateBody(Body b, Direction d, int speed) {
		GridCell	origCell = grid.getCell(b.getX(), b.getY()),
					currCell = origCell,
					targetCell;
		int i;
		for ( i = 0; i < speed; i++) {
			targetCell = currCell.go(d);
			if ( targetCell.isBorder() || !move(b, d, targetCell) )
				break;
			currCell = targetCell;
		}
		if ( i > 0 ) {
			origCell.removeBody(b);
			currCell.addBody(b);
			b.setLocation(currCell.getCenterX(), currCell.getCenterY());
			return true;
		}
		return false;
	}
	
	protected boolean move( Body b, Direction d, GridCell targetCell ) {
		//TODO Don't like all these fors and ifs
		List<Body> occupiers = targetCell.getOccupiers();
		for ( Body targetBody : occupiers )
			if ( targetBody.getType() == Body.OBSTACLE ) {
				fireCollisionEvent(b, targetBody, new Point2D.Double(d.x(), d.y()));
				return false;
			}
		for ( Body targetBody : occupiers )
			if ( targetBody.getType() == Body.MOVABLE ) {
				fireCollisionEvent(b, targetBody, new Point2D.Double(d.x(), d.y()));
				if ( !translateBody(targetBody, d, 1) )
					return false;
			}
		//just create collisionevent for the other body types
		for ( Body targetBody : occupiers )
			if ( targetBody.getType() != Body.OBSTACLE && targetBody.getType() != Body.MOVABLE ) {
				fireCollisionEvent(b, targetBody, new Point2D.Double(d.x(), d.y()));
				if ( targetBody.getType() == Body.GOAL )
					episodeEnd = true;
			}
				
		return true;
	}

	@Override
	public Grid getGrid() {
		return grid;
	}
	
	
	public RewardFunction getRewardFunction() {
		return reward;
	}

	public void setRewardFunction(RewardFunction reward) {
		this.reward = reward;
	}
	
	protected void addToCells(Body b) {
    	Rectangle2D bounds = b.getFigure().getBounds2D();
    	if ( bounds.getWidth() > grid.getCellWidth() || bounds.getHeight() > grid.getCellHeight() ) {
    		GridCell[][] cells = grid.getCells();
    		for(int h = 0; h < cells[0].length; h++)
        		for(int w = 0; w < cells.length; w++)
        			if ( cells[w][h].intersects(bounds) )
        				cells[w][h].addBody(b);
    	}
    	else {
    		Point2D pos = b.getLocation();
    		GridCell cell = grid.getCell(pos.getX(), pos.getY());
    		cell.addBody(b);
    		if ( !cell.contains(bounds) ) {
    			GridCell[] neighbours = cell.getQuadrantNeightbours(pos.getX(), pos.getY());
    			for ( int i = 0; i < neighbours.length; i++ )
    				if ( neighbours[i].intersects(bounds) )
    					neighbours[i].addBody(b);
    		}    		
    	}
    }
	
	protected void removeFromCells(Body b) {
    	Rectangle2D bounds = b.getFigure().getBounds2D();
    	if ( bounds.getWidth() > grid.getCellWidth() || bounds.getHeight() > grid.getCellHeight() ) {
    		GridCell[][] cells = grid.getCells();
    		for(int h = 0; h < cells[0].length; h++)
        		for(int w = 0; w < cells.length; w++)
        			if ( cells[w][h].intersects(bounds) )
        				cells[w][h].removeBody((net.javlov.world.Body)b);
    	}
    	else {
    		Point2D pos = b.getLocation();
    		GridCell cell = grid.getCell(pos.getX(), pos.getY());
    		cell.removeBody(b);
    		if ( !cell.contains(bounds) ) {
    			GridCell[] neighbours = cell.getQuadrantNeightbours(pos.getX(), pos.getY());
    			for ( int i = 0; i < neighbours.length; i++ )
    				if ( neighbours[i].intersects(bounds) )
    					neighbours[i].removeBody(b);
    		}    		
    	}
    }
	
	protected void fireCollisionEvent(Body b1, Body b2, Point2D.Double speed) {
		CollisionEvent event = new CollisionEvent(b1, b2, speed, (Point2D.Double)b2.getLocation());
		for ( CollisionListener listener : listeners )
			listener.collisionOccurred(event);
	}
}
